package com.ag.improvedjava.iterable;

import java.util.Comparator;
import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.Set;
import java.util.TreeSet;

final class UnionIterable<T> implements Iterable<T> {

	private Iterable<T> iterable1;
	private Iterable<T> iterable2;
	private Comparator<T> comparer;
	
	UnionIterable(Iterable<T> iterable1, Iterable<T> iterable2) {
		this.iterable1 = iterable1;
		this.iterable2 = iterable2;
	}

	UnionIterable(Iterable<T> iterable1, Iterable<T> iterable2, Comparator<T> comparer) {
		this.iterable1 = iterable1;
		this.iterable2 = iterable2;
		this.comparer = comparer;
	}

	@Override
	public Iterator<T> iterator() {
		return new UnionIterator<>(iterable1.iterator(), iterable2.iterator(), comparer);
	}
}

class UnionIterator<T> extends UnmodifiableIterator<T> {
	
	private Iterator<T> iterator1;
	private Iterator<T> iterator2;
	private Set<T> distinctElements;
	private T current;
	
	private boolean canIterate = true;
	private boolean hasNext = false;
	
	UnionIterator(Iterator<T> iterator1, Iterator<T> iterator2, Comparator<T> comparer) {
		this.iterator1 = iterator1;
		this.iterator2 = iterator2;
		this.distinctElements = new TreeSet<T>(comparer);
	}

	@Override
	public boolean hasNext() {
		while (canIterate && (iterator1.hasNext() || iterator2.hasNext())) {
			hasNext = false;
			current = iterator1.hasNext() ? iterator1.next() : iterator2.next();
			if (distinctElements.add(current)) {
				canIterate = false;
				hasNext = true;
			}
		}
		if (canIterate) {
			hasNext = false;
		}
		return hasNext;
	}

	@Override
	public T next() {
		if (!hasNext()) {
			throw new NoSuchElementException();
		}
		canIterate = true;
		return current;
	}

}
